#include "stdafx.h"
#include <iostream>
#include <cmath>
#include <opencv2\opencv.hpp>
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <opencv2/core/eigen.hpp>
#include <vector>

#include "time.h"


using namespace Eigen;

using namespace cv;
using namespace std;

enum ConvolutionType {   
	/* Return the full convolution, including border */
	CONVOLUTION_FULL, 

	/* Return only the part that corresponds to the original image */
	CONVOLUTION_SAME,

	/* Return only the submatrix containing elements that were not influenced by the border */
	CONVOLUTION_VALID
};

// utilities
void dumpMat(const MatrixXf &in)
{
	for(int i=0;i<in.rows();i++)
	{
		for(int j=0;j<in.cols();j++)
		{
			printf("%1.4f\t\t",in(i,j));
		}
		cout << endl;
	}
	cout << endl;
}
void dumpCVMat(const cv::Mat &in)
{
	for(int i=0;i<in.rows;i++)
	{
		for(int j=0;j<in.cols;j++)
		{
			printf("%1.4f\t\t",in.at<float>(i,j));
		}
		cout << endl;
	}
	cout << endl;
}

Mat diff(const Mat &x,const int direction)
{
	// backward difference, padded to the same size of x ( padarray-free version )
	// direction: 0 -> x
	//            1 -> y
	int c=x.channels();
	if(c==3)
	{
		if(direction==0)
		{
			Mat result(x.size(),CV_32FC3);
			result=Mat::zeros(x.rows,x.cols,CV_32FC3);
			for(int i=0;i<x.rows;i++)
			{
				for(int j=1;j<x.cols;j++)
				{
					for(int c=0;c<3;c++)
					{
						// each channel
						result.at<Vec3f>(i,j-1)[c]=x.at<Vec3f>(i,j)[c]-x.at<Vec3f>(i,j-1)[c];
					}
				}
			}
			return result;
		}
		else
		{
			Mat result(x.size(),CV_32FC3);
			result=Mat::zeros(x.rows,x.cols,CV_32FC3);
			for(int i=1;i<x.rows;i++)
			{
				for(int j=0;j<x.cols;j++)
				{
					for(int c=0;c<3;c++)
					{
						// each channel
						result.at<Vec3f>(i-1,j)[c]=x.at<Vec3f>(i,j)[c]-x.at<Vec3f>(i-1,j)[c];
					}
				}
			}
			return result;
		}
	}
	else
	{
		// single channel image
		if(direction==0)
		{
			Mat result(x.size(),CV_32FC1);
			result=Mat::zeros(x.size(),CV_32FC1);
			for(int i=0;i<x.rows;i++)
			{
				for(int j=1;j<x.cols;j++)
				{
					result.at<float>(i,j-1)=x.at<float>(i,j)-x.at<float>(i,j-1);
				}
			}
			return result;
		}
		else
		{
			Mat result(x.size(),CV_32FC1);
			result = Mat::zeros(x.size(),CV_32FC1);
			for(int i=1;i<x.rows;i++)
			{
				for(int j=0;j<x.cols;j++)
				{
						result.at<float>(i-1,j)=x.at<float>(i,j)-x.at<float>(i-1,j);
				}
			}
			return result;
		}
	}
}

Mat flip(const Mat &fin)
{
	Mat result;
	cv::flip(fin,result,-1);
	return result;
}


void conv2(const Mat &img, const Mat& kernel, ConvolutionType type, Mat& dest) 
{
	// simulate Matlab conv2 behavior.
	Mat source = img;
	if(CONVOLUTION_FULL == type) 
	{
		source = Mat();
		const int additionalRows = kernel.rows-1, additionalCols = kernel.cols-1;
		copyMakeBorder(img, source, (additionalRows+1)/2, additionalRows/2, (additionalCols+1)/2, additionalCols/2, BORDER_CONSTANT, Scalar(0));
	}

	Point anchor(kernel.cols - kernel.cols/2 - 1, kernel.rows - kernel.rows/2 - 1);
	int borderMode = BORDER_CONSTANT;
	filter2D(source, dest, img.depth(), flip(kernel), anchor, 0, borderMode);

	if(CONVOLUTION_VALID == type) 
	{
		dest = dest.colRange((kernel.cols-1)/2, dest.cols - kernel.cols/2)
			.rowRange((kernel.rows-1)/2, dest.rows - kernel.rows/2);
	}
}

Mat lpfilter(const Mat &fin, const double sigma)
{
	// Gaussian filter.
	// A really bad news is the implementation difference.
	// Convolution might be optimized. But is is fast enough now.
	
	int ksize=floor(5*sigma+0.5)>1?floor(5*sigma+0.5):1;
	if(ksize%2==0)
		ksize+=1;
	Mat result;
	Mat kernel=getGaussianKernel(ksize, sigma,CV_32F);
	
	conv2(fin,kernel,CONVOLUTION_SAME,result);
	
	transpose(kernel,kernel);
	conv2(result,kernel,CONVOLUTION_SAME,result);

	return result;
	
}
void computeTextureWeights(const Mat &fin, double sigma,double sharpness, vector<Mat> &retWeights)
{
	Mat fx = diff(fin,0);
	Mat fy = diff(fin,1);

	double vareps_s = sharpness;
	double vareps = 0.001;

	Mat wto;
	double c=fx.channels();
	if(c==1)
	{
		Mat wt(fin.size(),CV_32FC1);
		wt = Mat::zeros(fin.size(),CV_32FC1);
		for(int x=0;x<fx.rows;x++)
		{
			for(int y=0;y<fx.cols;y++)
			{
				double l=fx.at<float>(x,y);
				double r=fy.at<float>(x,y);
				wt.at<float>(x,y)=1/(std::max(sqrt(l*l+r*r),vareps_s));
			}
		}
		wt.copyTo(wto);
	}
	else
	{
		// assume c==3;
		Mat wt(fin.size(),CV_32FC1);
		wt = Mat::zeros(fin.size(),CV_32FC1);
		for(int x=0;x<fx.rows;x++)
		{
			for(int y=0;y<fx.cols;y++)
			{
				double mag=0;
				for(int t=0;t<3;t++)
				{
					double l=fx.at<Vec3f>(x,y)[t];
					double r=fy.at<Vec3f>(x,y)[t];
					mag+=sqrt(l*l+r*r);
				}
				wt.at<float>(x,y)=1/(std::max(mag/3,vareps_s));
			}
		}
		wt.copyTo(wto);
	}

	
	Mat fbin = lpfilter(fin, sigma);
	Mat gfx = diff(fbin,0);
	Mat gfy = diff(fbin,1);

	vector <Mat> mv;
	split(fbin,mv);
	
	Mat wtbx,wtby;
	if(c==1)
	{
		Mat wtx=Mat::zeros(fin.size(),CV_32FC1);
		Mat wty=Mat::zeros(fin.size(),CV_32FC1);
		for(int x=0;x<gfx.rows;x++)
		{
			for(int y=0;y<gfx.cols;y++)
			{
				double l=gfx.at<float>(x,y);
				double r=gfy.at<float>(x,y);
				wtx.at<float>(x,y)=1/(std::max(fabs(l),vareps));
				wty.at<float>(x,y)=1/(std::max(fabs(r),vareps));
			}
		}
		wtx.copyTo(wtbx);
		wty.copyTo(wtby);
	}
	else
	{
		// assume c==3;
		Mat wtx=Mat::zeros(fin.size(),CV_32FC1);
		Mat wty=Mat::zeros(fin.size(),CV_32FC1);
		for(int x=0;x<gfx.rows;x++)
		{
			for(int y=0;y<gfx.cols;y++)
			{
				double magX=0;
				double magY=0;
				for(int t=0;t<3;t++)
				{
					double l=gfx.at<Vec3f>(x,y)[t];
					double r=gfy.at<Vec3f>(x,y)[t];
					magX+=fabs(l);
					magY+=fabs(r);
				}
				wtx.at<float>(x,y)=0.001/(std::max(magX/3,vareps));
				wty.at<float>(x,y)=0.001/(std::max(magY/3,vareps));
			}
		}
		wtx.copyTo(wtbx);
		wty.copyTo(wtby);
	}
	
	Mat retx(fin.size(),CV_32FC1);
	Mat rety(fin.size(),CV_32FC1);

	retx=wtbx.mul(wto);
	rety=wtby.mul(wto);

	for(int x=0;x<retx.cols;x++)
	{
		retx.at<float>(retx.rows-1,x)=0;
		rety.at<float>(rety.rows-1,x)=0;
	}
	for(int y=0;y<retx.rows;y++)
	{
		retx.at<float>(y,retx.cols-1)=0;
		rety.at<float>(y,rety.cols-1)=0;
	}

	retWeights.push_back(retx);
	retWeights.push_back(rety);
}
SparseMatrix<float> spdiags(const MatrixXf &a_in, const vector<int> lines, const int dimx, const int dimy)
{
	// insert the diagonal data for sparse matrices.
	SparseMatrix<float> result(dimx,dimy);
	for(int i=0;i<lines.size();i++)
	{
		int count=dimx-abs(lines[i]);
		if(lines[i]<=0)
		{
			// assumption: dimx==dimy
			// in this case, it works ( parameters are always k*k ).
			// might be fixed for future use.
			for(int j=0;j<count;j++)
			{
				result.insert(j-lines[i],j)=a_in(j,i);
			}
		}
		if(lines[i]>0)
		{
			for(int j=0;j<count;j++)
			{
				result.insert(j,j+lines[i])=a_in(j+lines[i],i);
			}
		}
	}
	return result;
}

MatrixXf padarray(const MatrixXf &a_in, const int count, const int type)
{
	// type :  0 -> 'pre'
	// in fact, the padded array is a vector.

	MatrixXf result (a_in.rows()+count,a_in.cols());
	result.fill(0);
	for (int x=0;x<a_in.cols();x++)
	{
		for(int y=0;y<a_in.rows();y++)
		{
			result(y+count,x)=a_in(y,x);
		}
	}
	return result;
}

MatrixXf cvToEigen(const cv::Mat &in)
{
	MatrixXf result(in.rows,in.cols);
	for(int i=0;i<in.rows;i++)
		for(int j=0;j<in.cols;j++)
			result(i,j)=in.at<float>(i,j);
	return result;
}

Mat solveLinearEquation(const Mat &m_IN,const Mat &wx,const Mat &wy, const double lambda)
{
	MatrixXf e_IN,e_wx,e_wy;
	
	// cv2eigen and eigen2cv are OpenCV built-in functions.
	// Used for fast data copy from CV to Eigen.
	cv2eigen(m_IN,e_IN);
	cv2eigen(wx,e_wx);
	cv2eigen(wy,e_wy);
	
	int r = m_IN.rows;
	int c = m_IN.cols;
	int ch = m_IN.channels();
	int ks = r*c;

	MatrixXf e_dx = e_wx;
	MatrixXf e_dy = e_wy;
	e_dx.resize(ks,1);
	e_dx = e_dx*-lambda;
	e_dy.resize(ks,1);
	e_dy = e_dy*-lambda;
	
	MatrixXf e_B(ks,2);
	e_B.col(0)=e_dx.col(0);
	e_B.col(1)=e_dy.col(0);
	
	vector <int> d;
	d.push_back(-r);
	d.push_back(-1);
	
	SparseMatrix<float> e_A = spdiags(e_B,d,ks,ks);
	MatrixXf e_e = e_dx;
	MatrixXf e_w = padarray(e_dx, r, 0); 
	e_w.topRows(e_w.rows()-r);

	MatrixXf e_s = e_dy;
	MatrixXf e_n = padarray(e_dy, 1, 0);
	e_n.topRows(e_n.rows()-1);

	MatrixXf ones(e_e.rows(),e_e.cols());
	ones.fill(0.001);

	MatrixXf e_D = ones-(e_e+e_w+e_s+e_n);
	d.clear();
	d.push_back(0);
	e_A = e_A + SparseMatrix<float>(e_A.transpose()) + spdiags(e_D, d, ks, ks); 

	vector <cv::Mat> mv,mv_out;
	for(int i=0;i<3;i++)
		mv_out.push_back(Mat());
	split(m_IN,mv);


	// cuda implementation
	/*SparseMatrix<float> e_At = e_A.transpose();

	MatrixXf tmat;
	cv2eigen(mv[0],tmat);
	VectorXf tb,res;
	tmat.resize(ks,1);
	tb=tmat.col(0);

	
	int M = 0, N = 0, nz = 0, *I = NULL, *J = NULL;
	float *val = NULL;
	float *x,*rhs;

	M = N = e_A.rows();
	nz = e_At.nonZeros();
	I = (int *)malloc(sizeof(int)*(N+1));
	J = (int *)malloc(sizeof(int)*nz);
	val = (float *)malloc(sizeof(float)*nz);
	
	int lastrow = -1;
	int counter = 0;
	for (int k=0; k<e_At.outerSize(); ++k)
	{
		for (SparseMatrix<float>::InnerIterator it(e_At,k); it; ++it)
		{
			int row = it.col();
			int col = it.row();
			float v = it.value();

			val[counter] = v;
			J[counter] = col;
			if(row > lastrow)
			{
				lastrow +=1;
				I[lastrow] = counter;
			}
			counter += 1;
		}
	}
	I[N] = nz;

	x = (float *)malloc(sizeof(float)*N);
	rhs = (float *)malloc(sizeof(float)*N);

	for (int i = 0; i < N; i++)
	{
		rhs[i] = tb(i);
		x[i] = 0.0;
	}

	Mat result(r,c,CV_32FC3);
	
	int to = 0;
	float tol = pow(0.1,to);
	conjugateGradient(M,N,nz,I,J,val,x,rhs,tol);

	Mat t2(r,c,CV_8U);
	cout << r << " " << c << endl;
	for(int xt=0;xt<r;xt++)
	{
		for(int yt=0;yt<c;yt++)
		{
			t2.at<uchar>(xt,yt) = x[yt*r+xt]*0.255;
			for(int cc=0;cc<3;cc++)
				result.at<Vec3f>(xt,yt)[cc] = x[yt*r+xt]*0.001;
		}
		t2.at<uchar>(xt,c-1) = t2.at<uchar>(xt,c-2);
		for(int cc=0;cc<3;cc++)
			result.at<Vec3f>(xt,c-1)[cc] = result.at<Vec3f>(xt,c-2)[cc];
	}
	char fn[500];
	sprintf_s(fn,"t2_%d.png",to);
	//imwrite(fn,t2);
	
	free(x);*/

	

	for(int i=0;i<=0;i++)
	{
		// solve linear equation for each channel.
		MatrixXf tmat;
		cv2eigen(mv[i],tmat);
		VectorXf tb,res;
		tmat.resize(ks,1);
		tb=tmat.col(0);

		// CG Method. Set iterations to balance speed and quality.
		ConjugateGradient<SparseMatrix<float> > cg;
		cg.setMaxIterations(120);
		cg.compute(e_A);
		res=cg.solve(tb);
		MatrixXf tres = res;
		tres.resize(r,c);

		Mat t;
		eigen2cv(tres,t);
		t.convertTo(t,CV_32F,0.001,0);
		mv_out[i]=t;
	}
	
	return mv_out[0];
}
Mat tsmooth(const Mat &I,double lambda,double sigma,double sharpness, int maxIter)
{
	// set default parameters.
	if(lambda==0)
		lambda=0.01;
	if(sigma==0)
		sigma=3.0;
	if(sharpness==0)
		sharpness=0.02;
	if(maxIter==0)
		maxIter=5;
	std::cerr <<  I.channels() << endl;


	Mat dI(I.size(),CV_32FC3);
	I.convertTo(dI, CV_32FC3, 1/255.0);
	
	Mat x (dI);
	double sigma_iter = sigma;
	lambda = lambda/2.0;
	double dec=2.0;

	// Iteration
	for(int i=0;i<maxIter;i++)
	{
		vector <Mat> ret;
		clock_t t1,t2;
		t1=clock();
		computeTextureWeights(x, sigma_iter, sharpness, ret);
		t2=clock();
		float diff=((float)t2-(float)t1);
		cerr<<diff/CLOCKS_PER_SEC<<endl;
	
		Mat wx=ret[0];
		Mat wy=ret[1];
		t1=clock();
		x = solveLinearEquation(dI, wx, wy, lambda);
		t2=clock();
		diff=((float)t2-(float)t1);
		cerr<<diff/CLOCKS_PER_SEC<<endl;
		sigma_iter = sigma_iter/dec;
		if (sigma_iter < 0.5)
			sigma_iter = 0.5;
	}
	Mat t(x.size(),CV_8UC1);
	x.convertTo(t,CV_8UC1,255.0);
	imwrite("t.png",t);
	return t;
}